#include <stdio.h>
#include <stdlib.h>
#include <malloc.h>
#include <math.h>
#include <string.h>
#include <mpi.h>
#include <stdbool.h>

/*
 * Подключаем заголовочный файл с параметрами расчета.
 */
#include "config.h"

/*
 * Индексация по узлам сетки.
 * Ожнозначно определяет маппинг двумерной области индексов i, j на одномерную.
 */
#define ind(i, j) ((i) + (j) * t->n.x)

#define max(a, b) ((a) > (b) ? (a) : (b))
#define min(a, b) ((a) < (b) ? (a) : (b))
#define abs(a) ((a) < 0 ? -(a) : (a))

/*
 * Структура двумерного вектора.
 */
typedef struct {
	double x, y;
} vec2;

/*
 * Структура узла с данными и одновременно инварианта Римана.
 * Хранит в себе вектор скорости и давление, либо компоненты инварианта Римана wi.
 */
typedef union {
	struct { vec2 v; double p; };
	struct { double w1, w2, w3; };
} node_t;

typedef node_t w_t;

/*
 * Структура материала среды.
 * Хранит в себе плотность и скорость звука.
 */
typedef struct {
	double rho;
	double c;
} material_t;

/*
 * Структура целочисленной точки.
 * Используетс для индексации и работы с двумерными массивами.
 */
typedef struct {
	int x, y;
} point_t;

/*
 * Структура всей задачи.
 * Хранит в себе все необходимое для расчетов.
 */
typedef struct {
    
    /* номер процесса */
    int rank;

    /* кол-во процессов */
    int size;

    /* начало куска нашего процесса */
    int begin_y;

    /* высота полосы для данного процесса */
    int size_y;

	int* begin_y_arr;
	int* size_y_arr;

	/* Шаг интегрирования по времени. */
	double dt;
	
	/*
	 * Коэффициенты в схеме расщепления,
	 * a.x и a.y - соответственно по различным осям.
	 */
	vec2 a;
	
	/* Шаг расчетной сетки по соотв. осям.*/
	vec2 h;
	
	/* 
	 * Физическое положение сетки в пространстве.
	 * Координаты левого нижнего угла.
	 */
	vec2 origin;
	
	/* Размер сетки. Число узлов по осям. */
	point_t n;
	
	/*
	 * Положение каверны в расчетной сетке.
	 * Индексы узлов левого нижнего края каверны.
	 */
	point_t co;
	
	/* Размеры каверны в узлах разностной сетки. */
	point_t cn;
	
	/*
	 * Разностные сетки.
	 * 0 - значение на слое n, 1 - на слое n - 1.
	 */
	node_t *g[2];
	
	/* Материал среды. */
	material_t mat;
} task_t;

/*
 * Форма импульса волны возмущения. Берется из конфигурационного файла.
 */
double impulse(const double x) { return INIT_PULSE(x); }

/*
 * Проинициализировать структуру с задачей.
 * Заполнение всей полей структуры и выделение памяти.
 */
void init_task(task_t *t);

/*
 * Корректное удаление структуры задачи.
 * Очищает занимаемую ранее память.
 */
void free_task(task_t *t);

/*
 * Сохраняет задачу в файл в формате Legacy VTK.
 */
int  save_task(const char *file, task_t *t);

/*
 * Сохраняет текущее состояние сейсмодатчиков на поверхности расчетной области.
 * Если параметр first выставлен в 1, то пересоздает файл с сейсмограммой и начинает запись сначала.
 */
int  save_seismo(const char *file, task_t *t, int first);

/*
 * Утильная функция.
 * Меняет порядок байт в float.
 * Необходимо для корректной работы с бинарным форматом VTK.
 */
void write_float(FILE *f, const float v);

/*
 * Выставить начальные условия сферического взрыва.
 */
void init_ball(task_t *t, const double r1, const double r2, const vec2 o, const double mag);

/*
 * Выставить начальные условия фронта проской продолной волны.
 */
void init_wave(task_t *t, const double w, const vec2 o, const vec2 dir, const double mag);

/*
 * Вспомогательные функции для работы с векторами.
 */
vec2 vadd(const vec2 a, const vec2 b) { vec2 t = { a.x + b.x, a.y + b.y }; return t; }
vec2 vsub(const vec2 a, const vec2 b) { vec2 t = { a.x - b.x, a.y - b.y }; return t; }
double vdot(const vec2 a, const vec2 b) { return a.x * b.x + a.y * b.y; }
double vnorm(const vec2 a) { return sqrt(vdot(a, a)); }
vec2 vscal(const vec2 a, const double s) { vec2 t = { a.x * s, a.y * s }; return t; }

/*
 * Делает домножение узла на матрицу омега большое для координаты x.
 * w = Omega_x * n
 */
void omega_x(const node_t n, const material_t m, w_t *w)
{
	double k = n.p / (m.c * m.rho);
	w->w1 = n.v.x + k;
	w->w2 = n.v.x - k;
}

/*
 * Делает домножение узла на матрицу омега большое для координаты y.
 * w = Omega_y * n
 */
void omega_y(const node_t n, const material_t m, w_t *w)
{
	double k = n.p / (m.c * m.rho);
	w->w1 = n.v.y + k;
	w->w2 = n.v.y - k;
}

/*
 * Делает домножение узла на матрицу обратную омега большое для координаты x.
 */
void inc_x(const w_t d, const material_t m, node_t *n)
{
	n->v.x += 0.5 * (d.w1 + d.w2);
	n->p   += 0.5 * m.c * m.rho * (d.w1 - d.w2);
}

/*
 * Делает домножение узла на матрицу обратную омега большое для координаты y.
 */
void inc_y(const w_t d, const material_t m, node_t *n)
{
	n->v.y += 0.5 * (d.w1 + d.w2);
	n->p   += 0.5 * m.c * m.rho * (d.w1 - d.w2);
}

/*
 * Сеточно-характеристический ограничитель.
 */
double minmax(const double u, const double u1, const double u2)
{
	double vmin = min(u1, u2);
	double vmax = max(u1, u2);
	double v = u;
	if (v > vmax) v = vmax;
	if (v < vmin) v = vmin;
	return v;
}

/*
 * Схема Куранта-Изаксона-Риса первого порядка точности.
 */
double cir(const double k, const double u_1, const double u)
{
	return u - k * (u - u_1);
}

/*
 * Корректировка граничного условия свободной границы.
 */
void correct(node_t *n, vec2 norm, material_t mat)
{
	vec2 corr = vscal(norm, n->p / mat.rho / mat.c / vnorm(norm));
	n->v = vadd(n->v, corr);
	n->p = 0.0;
}

/*
 * Схема Русанова третьего порядка точности.
 */
double rusanov3(const double c, const double u1, const double u2, const double u3, const double u4)
{
	const double d0  = u2 - u3;
	const double d1  = u3 - u4;
	const double d_1 = u1 - u2;
	double v = u3 + 0.5 * c * (d0 + d1) + 0.5 * c * c * (d0 - d1) + c * c * (c - 1) / 6.0 * (d_1 - 2.0 * d0 + d1);
	return minmax(v, u2, u3);
}

/*
 * Схема третьего порядка на смещенном шаблоне.
 */
double third1(const double c, const double u1, const double u2, const double u3, const double u4)
{
	const double d0 = u1 - u2;
	const double d1 = u2 - u3;
	const double d2 = u3 - u4;
	double v = u2 + 0.5 * c * (d0 + d1) + 0.5 * c * c * (d0 - d1) + c * c * (c - 1) / 6.0 * (d0 - 2.0 * d1 + d2);
	return minmax(v, u1, u2);
}

/*
 * Выполнить шаг расщепления в направлении x для заданной области сетки.
 */
void step_x3(task_t *t, point_t p1, point_t p2)
{
	// Переменные для цикла.
	int i, j;
	
	/*
	 * Изменяем шаг по времени исходя из шага расщепления.
	 * tau_x = tau / alpha_x
	 */
	double tau = t->dt / t->a.x;
	
	// Число Куранта
	double k = tau * t->mat.c / t->h.x;
	
	// По всем координатам по y.
	for (j = p1.y; j < p2.y; j++) {
		/*
		 * Структуры для хранения инвариантов Римана по нескольким индексам координаты.
		 * Число структур определяется шаблоном разностной схемы.
		 * У нас все шаблоны пятиточечные.
		 * Соответственно w[m-2], w[m-1], w[m], w[m+1], w[m+2],
		 * m - номер текущего узла по x.
		 */
		w_t w_2, w_1, w, w1, w2;
		
		// Делаем преобразование первых узлов.
		omega_x(t->g[1][ind(p1.x,     j)], t->mat, &w);
		omega_x(t->g[1][ind(p1.x + 1, j)], t->mat, &w1);
		omega_x(t->g[1][ind(p1.x + 2, j)], t->mat, &w2);
		w_2 = w_1 = w;
		
		// Цикл по координате x.
		for (i = p1.x; i < p2.x; i++) {
			// Приращение инварианта Римана за данный шаг.
			w_t d;
			
			/*
			 * Если слева достаточно ячеек, то делаем шаг по схеме Русанова.
			 *     +
			 *     \
			 * +-+-+-+
			 */
			if (i >= p1.x + 2 && i < p2.x - 1) {
				d.w1 = rusanov3(k, w_2.w1, w_1.w1, w.w1, w1.w1) - w.w1;
			}
			/*
			 * Если ячеек не достаточно делаем по схеме третьего порядка на смещенном шаблоне.
			 *   +
			 *   \
			 * +-+-+-+
			 */
			else if(i >= p1.x + 1 && i < p2.x - 2) {
				d.w1 = third1(k, w_1.w1, w.w1, w1.w1, w2.w1) - w.w1;
			} 
			/*
			 * В крайнем случае делаем уголком.
			 * +
			 * \
			 * +-+
			 */
			else {
				d.w1 = cir(k, w_1.w1, w.w1) - w.w1;
			}
			
			// Аналогично для другой характеристики, только в другую сторону.
			if (i >= p1.x && i < p2.x - 2) {
				d.w2 = rusanov3(k, w2.w2, w1.w2, w.w2, w_1.w2) - w.w2;
			} else if (i >= 2 && i < p2.x - 1) {
				d.w2 = third1(k, w1.w2, w.w2, w_1.w2, w_2.w2) - w.w2;
			} else {
				d.w2 = cir(k, w1.w2,  w.w2) - w.w2;
			}
			t->g[0][ind(i, j)] = t->g[1][ind(i, j)];
			
			// Обновляем значене в сетке.
			inc_x(d, t->mat, t->g[0] + ind(i, j));
			
			// Преобразуем следующие узлы.
			w_2 = w_1;
			w_1 = w;
			w = w1;
			w1 = w2;
			if (i < p2.x - 3) omega_x(t->g[1][ind(i + 3, j)], t->mat, &w2);
		}
	}
}

/*
 * Выполнить шаг расщепления в направлении y для заданной области сетки.
 */

bool in_cavern(task_t* t, int i, int j)
{
    if(t->co.x <= i && i <= t->co.x + t->cn.x && //in cavern by x-coord
       t->co.y <= j && j <= t->co.y + t->cn.y)
        return true;
    return false;
}

void step_y3(task_t *t, point_t p1, point_t p2)
{
	int i, j;
	double tau = t->dt / t->a.y;
	double k = tau * t->mat.c / t->h.y;
	for (i = p1.x; i < p2.x; i++) {
		w_t w_2, w_1, w, w1, w2;
		omega_y(t->g[1][ind(i, p1.y)],     t->mat, &w);
		omega_y(t->g[1][ind(i, p1.y + 1)], t->mat, &w1);
		omega_y(t->g[1][ind(i, p1.y + 2)], t->mat, &w2);
		w_2 = w_1 = w;
        if(p1.y > 0 && !in_cavern(t, i, p1.y - 1)) //TEST CAVERN BEHAVE ME
            omega_y(t->g[1][ind(i, p1.y - 1)], t->mat, &w_1);
        if(p1.y > 1 && !in_cavern(t, i, p1.y - 2))
            omega_y(t->g[1][ind(i, p1.y - 2)], t->mat, &w_2);

		for (j = p1.y; j < p2.y; j++) {
			w_t d;
			if (j >= p1.y + 2 && j < p2.y - 1) {
				d.w1 = rusanov3(k, w_2.w1, w_1.w1, w.w1, w1.w1) - w.w1;
			} else if(j >= p1.y + 1 && j < p2.y - 2) {
				d.w1 = third1(k, w_1.w1, w.w1, w1.w1, w2.w1) - w.w1;
			} else {
				d.w1 = cir(k, w_1.w1, w.w1) - w.w1;
			}
			if (j >= p1.y && j < p2.y - 2) {
				d.w2 = rusanov3(k, w2.w2, w1.w2, w.w2, w_1.w2) - w.w2;
			} else if (j >= 2 && j < p2.y - 1) {
				d.w2 = third1(k, w1.w2, w.w2, w_1.w2, w_2.w2) - w.w2;
			} else {
				d.w2 = cir(k, w1.w2,  w.w2) - w.w2;
			}
			t->g[0][ind(i, j)] = t->g[1][ind(i, j)];
			inc_y(d, t->mat, t->g[0] + ind(i, j));
			w_2 = w_1;
			w_1 = w;
			w = w1;
			w1 = w2;
            
            if(!in_cavern(t, i, j + 3) && (j + 3 < t->n.y))
                omega_y(t->g[1][ind(i, j + 3)], t->mat, &w2);
		}
	}
}

/*
 * Выполнить шаг расщепления в направлении x для всей области сетки,
 * с учетом наличия в области каверны.
 * Также производится корректировка граничных условий.
 */
void step_x(task_t *t)
{
    point_t p1, p2;
    //ccccccccccccccccccccc    ---------------------
    //ccccccccccccccccccccc    ---------------------
    //---------------------    ---------------------
    //-------*****---------    -------*****---------
    //-------*****---------    -------*****---------
    //-------*****---------    -------*****---------
    //-------*****---------    -------*****---------
    //---------------------    ccccccccccccccccccccc
    //---------------------    ccccccccccccccccccccc
    //---------------------    ---------------------
    if(t->begin_y + t->size_y <= t->co.y           ||
       t->begin_y             >= t->co.y + t->cn.y)
    {
        p1.x = 0;      p1.y = t->begin_y;
        p2.x = t->n.x; p2.y = t->begin_y + t->size_y;
        step_x3(t, p1, p2);
    }
    //---------------------
    //---------------------
    //ccccccccccccccccccccc
    //ccccccc*****ccccccccc
    //-------*****---------
    //-------*****---------
    //-------*****---------
    //---------------------
    //---------------------
    //---------------------
    else if(t->begin_y             <= t->co.y           && 
            t->begin_y + t->size_y >= t->co.y           &&
            t->begin_y + t->size_y <= t->co.y + t->cn.y)
    {
        p1.x = 0;      p1.y = t->begin_y;
        p2.x = t->n.x; p2.y = t->co.y;
        step_x3(t, p1, p2);

        p1.x = 0;       p1.y = p2.y;
        p2.x = t->co.x; p2.y = t->begin_y + t->size_y;
        step_x3(t, p1, p2);

        p1.x = p2.x + t->cn.x; p1.y = p1.y;
        p2.x = t->n.x;         p2.y = p2.y;
        step_x3(t, p1, p2);
        
        int i;
        vec2 n1 = {1.0,  0.0};
        vec2 n2 = {-1.0, 0.0};
        for (i = 0; i < p2.y - p1.y; i++)
        {
            correct(t->g[0] + ind(t->co.x - 1,       t->co.y + i), n1, t->mat);
            correct(t->g[0] + ind(t->co.x + t->cn.x, t->co.y + i), n2, t->mat);
        }
    }
    //---------------------
    //---------------------
    //---------------------
    //-------*****---------
    //ccccccc*****ccccccccc
    //ccccccc*****ccccccccc
    //-------*****---------
    //---------------------
    //---------------------
    //---------------------
    else if(t->begin_y             >= t->co.y &&
            t->begin_y + t->size_y <= t->co.y + t->cn.y)
    {
        p1.x = 0;       p1.y = t->begin_y;
        p2.x = t->co.x; p2.y = t->begin_y + t->size_y;
        step_x3(t, p1, p2);

        p1.x = p2.x + t->cn.x; p1.y = p1.y;
        p2.x = t->n.x;         p2.y = p2.y;
        step_x3(t, p1, p2);
        
        int i;
        vec2 n1 = {1.0,  0.0};
        vec2 n2 = {-1.0, 0.0};
        for (i = 0; i < p2.y - p1.y; i++)
        {
            correct(t->g[0] + ind(t->co.x - 1,       t->begin_y + i), n1, t->mat);
            correct(t->g[0] + ind(t->co.x + t->cn.x, t->begin_y + i), n2, t->mat);
        }
    }
    //---------------------
    //---------------------
    //---------------------
    //-------*****---------
    //-------*****---------
    //-------*****---------
    //ccccccc*****ccccccccc
    //ccccccccccccccccccccc
    //---------------------
    //---------------------
    else if(t->begin_y             >= t->co.y           &&
            t->begin_y             <= t->co.y + t->cn.y &&
            t->begin_y + t->size_y >= t->co.y + t->cn.y)
    {
        p1.x = 0;      p1.y = t->co.y;
        p2.x = t->n.x; p2.y = t->begin_y + t->size_y;
        step_x3(t, p1, p2);

        p1.x = 0;       p1.y = t->begin_y;
        p2.x = t->co.x; p2.y = t->co.y;
        step_x3(t, p1, p2);

        p1.x = p2.x + t->cn.x; p1.y = p1.y;
        p2.x = t->n.x;         p2.y = p2.y;
        step_x3(t, p1, p2);
        
        int i;
        vec2 n1 = {1.0,  0.0};
        vec2 n2 = {-1.0, 0.0};
        for (i = 0; i < p2.y - p1.y; i++)
        {
            correct(t->g[0] + ind(t->co.x - 1,       t->co.y + i), n1, t->mat);
            correct(t->g[0] + ind(t->co.x + t->cn.x, t->co.y + i), n2, t->mat);
        }
    }
    //---------------------
    //---------------------
    //ccccccccccccccccccccc
    //ccccccc*****ccccccccc
    //ccccccc*****ccccccccc
    //ccccccc*****ccccccccc
    //ccccccc*****ccccccccc
    //ccccccccccccccccccccc
    //---------------------
    //---------------------
    else if(t->begin_y             <= t->co.y &&
            t->begin_y + t->size_y >= t->co.y + t->cn.y)
    {
        point_t p1, p2;
        p1.x = 0;      p1.y = t->begin_y;
        p2.x = t->n.x; p2.y = t->co.y;
        step_x3(t, p1, p2);

        p1.x = 0;      p1.y = t->co.y + t->cn.y;
        p2.x = t->n.x; p2.y = t->begin_y + t->size_y;
        step_x3(t, p1, p2);

        p1.x = 0;       p1.y = t->co.y;
        p2.x = t->co.x; p2.y = t->co.y + t->cn.y;
        step_x3(t, p1, p2);

        p1.x = p2.x + t->cn.x; p1.y = p1.y;
        p2.x = t->n.x;         p2.y = p2.y;
        step_x3(t, p1, p2);
        
        int i;
        vec2 n1 = {1.0,  0.0};
        vec2 n2 = {-1.0, 0.0};
        for (i = 0; i < t->cn.y; i++)
        {
            correct(t->g[0] + ind(t->co.x - 1,       t->co.y + i), n1, t->mat);
            correct(t->g[0] + ind(t->co.x + t->cn.x, t->co.y + i), n2, t->mat);
        }
    }
}

/*
 * Выполнить шаг расщепления в направлении y для всей области сетки,
 * с учетом наличия в области каверны.
 * Также производится корректировка граничных условий.
 */

void step_y(task_t *t)
{
    point_t p1, p2;
    //ccccccccccccccccccccc    ---------------------
    //ccccccccccccccccccccc    ---------------------
    //---------------------    ---------------------
    //-------*****---------    -------*****---------
    //-------*****---------    -------*****---------
    //-------*****---------    -------*****---------
    //-------*****---------    -------*****---------
    //---------------------    ccccccccccccccccccccc
    //---------------------    ccccccccccccccccccccc
    //---------------------    ---------------------
    if(t->begin_y + t->size_y <= t->co.y           ||
       t->begin_y             >= t->co.y + t->cn.y)
    {
        p1.x = 0;      p1.y = t->begin_y;
        p2.x = t->n.x; p2.y = t->begin_y + t->size_y;
        step_y3(t, p1, p2);
    }
    //---------------------
    //---------------------
    //ccccccccccccccccccccc
    //ccccccc*****ccccccccc
    //-------*****---------
    //-------*****---------
    //-------*****---------
    //---------------------
    //---------------------
    //---------------------
    else if(t->begin_y             <= t->co.y           && 
            t->begin_y + t->size_y >= t->co.y           &&
            t->begin_y + t->size_y <= t->co.y + t->cn.y)
    {
        p1.x = 0;      p1.y = t->begin_y;
        p2.x = t->n.x; p2.y = t->co.y;
        step_y3(t, p1, p2);

        p1.x = 0;       p1.y = p2.y;
        p2.x = t->co.x; p2.y = t->begin_y + t->size_y;
        step_y3(t, p1, p2);

        p1.x = p2.x + t->cn.x; p1.y = p1.y;
        p2.x = t->n.x;         p2.y = p2.y;
        step_y3(t, p1, p2);
    }
    //---------------------
    //---------------------
    //---------------------
    //-------*****---------
    //ccccccc*****ccccccccc
    //ccccccc*****ccccccccc
    //-------*****---------
    //---------------------
    //---------------------
    //---------------------
    else if(t->begin_y             >= t->co.y &&
            t->begin_y + t->size_y <= t->co.y + t->cn.y)
    {
        p1.x = 0;       p1.y = t->begin_y;
        p2.x = t->co.x; p2.y = t->begin_y + t->size_y;
        step_y3(t, p1, p2);

        p1.x = p2.x + t->cn.x; p1.y = p1.y;
        p2.x = t->n.x;         p2.y = p2.y;
        step_y3(t, p1, p2);
    }
    //---------------------
    //---------------------
    //---------------------
    //-------*****---------
    //-------*****---------
    //-------*****---------
    //ccccccc*****ccccccccc
    //ccccccccccccccccccccc
    //---------------------
    //---------------------
    else if(t->begin_y             >= t->co.y           &&
            t->begin_y             <= t->co.y + t->cn.y &&
            t->begin_y + t->size_y >= t->co.y + t->cn.y)
    {
        p1.x = 0;      p1.y = t->co.y;
        p2.x = t->n.x; p2.y = t->begin_y + t->size_y;
        step_y3(t, p1, p2);

        p1.x = 0;       p1.y = t->begin_y;
        p2.x = t->co.x; p2.y = t->co.y;
        step_y3(t, p1, p2);

        p1.x = p2.x + t->cn.x; p1.y = p1.y;
        p2.x = t->n.x;         p2.y = p2.y;
        step_y3(t, p1, p2);
    }
    //---------------------
    //---------------------
    //ccccccccccccccccccccc
    //ccccccc*****ccccccccc
    //ccccccc*****ccccccccc
    //ccccccc*****ccccccccc
    //ccccccc*****ccccccccc
    //ccccccccccccccccccccc
    //---------------------
    //---------------------
    else if(t->begin_y             <= t->co.y &&
            t->begin_y + t->size_y >= t->co.y + t->cn.y)
    {
        p1.x = 0;      p1.y = t->begin_y;
        p2.x = t->n.x; p2.y = t->co.y;
        step_y3(t, p1, p2);

        p1.x = 0;       p1.y = t->co.y + t->cn.y;
        p2.x = t->n.x; p2.y = t->begin_y + t->size_y;
        step_y3(t, p1, p2);

        p1.x = 0;       p1.y = t->co.y;
        p2.x = t->co.x; p2.y = t->co.y + t->cn.y;
        step_y3(t, p1, p2);

        p1.x = p2.x + t->cn.x; p1.y = p1.y;
        p2.x = t->n.x;         p2.y = p2.y;
        step_y3(t, p1, p2);
    }

	int i;
	vec2 n1 = {0.0, 1.0};
	vec2 n2 = {0.0, -1.0};
	// верх
	for (i = 0; i < t->n.x; i++) {
		correct(t->g[0] + ind(i, t->n.y - 1), n1, t->mat);
	}
	
	// каверна
	for (i = 0; i < t->cn.x; i++) {
		correct(t->g[0] + ind(t->co.x + i, t->co.y - 1),       n1, t->mat);
		correct(t->g[0] + ind(t->co.x + i, t->co.y + t->cn.y), n2, t->mat);
	}
}

void bound(task_t* t)
{
    if(t->rank % 2 == 0)
    {
        if(t->rank < t->size - 1)
        {
            MPI_Send(t->g[1] + (t->begin_y + t->size_y - 2) * t->n.x, 2 * 3 * t->n.x,
                     MPI_DOUBLE, t->rank + 1, 0, MPI_COMM_WORLD);
            MPI_Recv(t->g[1] + (t->begin_y + t->size_y) * t->n.x, 2 * 3 * t->n.x,
                     MPI_DOUBLE, t->rank + 1, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        }
        if(t->rank > 0)
        {
            MPI_Send(t->g[1] + t->begin_y * t->n.x, 2 * 3 * t->n.x,
                     MPI_DOUBLE, t->rank - 1, 0, MPI_COMM_WORLD);
            MPI_Recv(t->g[1] + (t->begin_y - 2) * t->n.x, 2 * 3 * t->n.x,
                     MPI_DOUBLE, t->rank - 1, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        }
    }
    else
    {
        if(t->rank > 0)
        {
            MPI_Recv(t->g[1] + (t->begin_y - 2) * t->n.x, 2 * 3 * t->n.x,
                     MPI_DOUBLE, t->rank - 1, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Send(t->g[1] + t->begin_y * t->n.x, 2 * 3 * t->n.x,
                     MPI_DOUBLE, t->rank - 1, 0, MPI_COMM_WORLD);
        }
        if(t->rank < t->size - 1)
        {
            MPI_Recv(t->g[1] + (t->begin_y + t->size_y) * t->n.x, 2 * 3 * t->n.x,
                     MPI_DOUBLE, t->rank + 1, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Send(t->g[1] + (t->begin_y + t->size_y - 2) * t->n.x, 2 * 3 * t->n.x,
                     MPI_DOUBLE, t->rank + 1, 0, MPI_COMM_WORLD);
        }
    }
}

void save_file(task_t* t, char* name, int i)
{
    int j = 0;
    if(!t->rank)
        for(j = 1; j < t->size; j++)
            MPI_Recv(t->g[0] + t->begin_y_arr[j] * t->n.x, 3 * t->size_y_arr[j] * t->n.x,
                     MPI_DOUBLE, j, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        
    else
        MPI_Send(t->g[0] + t->begin_y * t->n.x, 3 * t->size_y * t->n.x,
                 MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);

	MPI_Barrier(MPI_COMM_WORLD);
	if(!t->rank)
	{
	    sprintf(name, SAVE_DUMP_NAME, i);
		save_task(name, t);
		save_seismo(SAVE_SEISMO_NAME, t, i == 0);
	}
}

int main(int argc, char *argv[])
{
	// Переменные для циклов.
	int i, j, k;
	
	// Название выходного файла.
	char name[100];
	
	// Структура с задачей.
	task_t t;
	
	// Число шагов разностной схемы.
	int steps = TIME_STEPS;
	
	// Как часто сохранять.
	int save = SAVE_STEP;

    int size = 0;
    int rank = 0;

    MPI_Init(&argc, &argv);
	
    // Создаем структуру с задачей.
    init_task(&t);

    printf("Number of nodes = %d MyRank = %d\n"
           "size_y = %d, begin_y = %d\n\n", t.size, t.rank, t.size_y, t.begin_y);

	// МЕсто под временный слой для схемы расщепления.
	double *tg = (double*)malloc(2 * sizeof(node_t) * t.n.x * t.n.y);
	for (i = 0; i < steps; i++)
    {
		if(!t.rank)
			printf("step %d of %d\n", i, steps);
		
        if(i % save == 0)
	        save_file(&t, name, i);	
		// Шаг по оси X.
		step_x(&t);
		
		// Сохраняем что получилось (каждый процесс свой кусок)
		memcpy(tg + 3 * t.n.x * t.begin_y, t.g[0] + t.n.x * t.begin_y,
               sizeof(node_t) * t.n.x * t.size_y);
		
		// Шаг по оси Y.
		step_y(&t);
		
		// Схема расщепления.
		double *g0 = (double*)t.g[0];
		for (k = 3 * t.begin_y * t.n.x;
             k < 3 * (t.begin_y + t.size_y) * t.n.x;
             k++)
			g0[k] = g0[k] * t.a.y + tg[k] * t.a.x;

		// Копируем текущий шаг в предыдущий (n+1 в n).
		memcpy(t.g[1] + t.begin_y * t.n.x,
               t.g[0] + t.begin_y * t.n.x,
               sizeof(node_t) * t.n.x * t.size_y);

        bound(&t);
		
	}
	
	// Освобождение памяти.
	free_task(&t);
	free(tg);
    
    MPI_Finalize();

	return 0;
}

void init_task(task_t* t)
{
	// число узлов в области
	t->n.x = N_X;
	t->n.y = N_Y;
    
    MPI_Comm_size(MPI_COMM_WORLD, &(t->size));
    MPI_Comm_rank(MPI_COMM_WORLD, &(t->rank));
	
	t->begin_y_arr = calloc(t->size, sizeof(int));
	t->size_y_arr  = calloc(t->size, sizeof(int));
	int i = 0;
	for(i = 0; i < t->size; ++i)
	{
		t->size_y_arr[i] = N_Y / t->size;
        //последнему процессу весь остаток
		if(i == t->size - 1)
			t->size_y_arr[i] += N_Y % t->size;
		t->begin_y_arr[i] = N_Y / t->size * i;
	}

    t->size_y = t->size_y_arr[t->rank];
    t->begin_y = t->begin_y_arr[t->rank];
	
	// шаг сетки
	t->h.x = (double)SIZE_L / (double)N_X;
	t->h.y = (double)SIZE_H / (double)N_Y;
	
	// положение в пространстве
	t->origin.x = 0.0;
	t->origin.y = -t->h.y * t->n.y;
	
	// шаг по времени
	t->dt = TAU;
	
	// материал
	t->mat.c   = MAT_C;
	t->mat.rho = MAT_RHO;
	
	// коэф. расщепления
	t->a.x = 1.0 / (1.0 + t->h.x / t->h.y);
	t->a.y = 1.0 / (1.0 + t->h.y / t->h.x);
	printf("alpha = {%f, %f}\n", t->a.x, t->a.y);
	
	// размер каверны
	t->cn.x = CAVERN_L / t->h.x;
	t->cn.y = CAVERN_H / t->h.y;
	printf("cavern size = {%d, %d}\n", t->cn.x, t->cn.y);
	
	// положение каверны
	t->co.x = (CAVERN_A - t->origin.x) / t->h.x;
	t->co.y = (CAVERN_B - t->origin.y) / t->h.y;
	printf("cavern origin = {%d, %d}\n", t->co.x, t->co.y);
	
	// выделяем место под расчетные сетки
	t->g[0] = (node_t*)malloc(sizeof(node_t) * t->n.x * t->n.y);
	t->g[1] = (node_t*)malloc(sizeof(node_t) * t->n.x * t->n.y);
	
	if (INIT_TYPE == 1) {
		vec2 o = {BALL_O_X, BALL_O_Y};
		init_ball(t, BALL_R1, BALL_R2, o, INIT_MAGNITUDE);
	} else if (INIT_TYPE == 2) {
		vec2 o = {WAVE_O_X, WAVE_O_Y};
		vec2 d = {WAVE_DIR_X, WAVE_DIR_Y};
		init_wave(t, WAVE_W, o, d, INIT_MAGNITUDE);
	}
	
	memcpy(t->g[1], t->g[0], sizeof(node_t) * t->n.x * t->n.y);
}

void free_task(task_t* t)
{
	if (t->n.x * t->n.y > 0) {
		free(t->g[0]);
		free(t->g[1]);
	}

	free(t->begin_y_arr);
	free(t->size_y_arr);
	
}

void write_float(FILE* f, const float v)
{
	union {
		float f;
		unsigned char b[4];
	} dat1, dat2;
	dat1.f = v;
	dat2.b[0] = dat1.b[3];
	dat2.b[1] = dat1.b[2];
	dat2.b[2] = dat1.b[1];
	dat2.b[3] = dat1.b[0];
	fwrite(dat2.b, sizeof(unsigned char), 4, f);
}


int save_task(const char* file, task_t *t)
{
	int i, j;
	FILE *fp = fopen(file, "w");
	if (fp == NULL) {
		perror("Failed to open file");
		return 1;
	}
	fprintf(fp, "# vtk DataFile Version 3.0\n");
	fprintf(fp, "vtk output\n");
	fprintf(fp, "BINARY\n");
	fprintf(fp, "DATASET STRUCTURED_POINTS\n");
	fprintf(fp, "DIMENSIONS %d %d 1\n", t->n.x, t->n.y);
	fprintf(fp, "SPACING %f %f 0.0\n", t->h.x, t->h.y);
	fprintf(fp, "ORIGIN %f %f 0.0\n", t->origin.x, t->origin.y);
	fprintf(fp, "POINT_DATA %d\n", t->n.x * t->n.y);
	
	// давление
	fprintf(fp, "SCALARS p float 1\n");
	fprintf(fp, "LOOKUP_TABLE p_table\n");
	for (j = 0; j < t->n.y; j++) {
		for (i = 0; i < t->n.x; i++) {
			float v = (float)t->g[0][ind(i, j)].p;
			write_float(fp, v);
		}
	}
	
	// скорости
	fprintf(fp, "SCALARS vx float 1\n");
	fprintf(fp, "LOOKUP_TABLE vx_table\n");
	for (j = 0; j < t->n.y; j++) {
		for (i = 0; i < t->n.x; i++) {
			float v = (float)t->g[0][ind(i, j)].v.x;
			write_float(fp, v);
		}
	}
	fprintf(fp, "SCALARS vy float 1\n");
	fprintf(fp, "LOOKUP_TABLE vy_table\n");
	for (j = 0; j < t->n.y; j++) {
		for (i = 0; i < t->n.x; i++) {
			float v = (float)t->g[0][ind(i, j)].v.y;
			write_float(fp, v);
		}
	}
	fprintf(fp, "SCALARS vnorm float 1\n");
	fprintf(fp, "LOOKUP_TABLE vy_table\n");
	for (j = 0; j < t->n.y; j++) {
		for (i = 0; i < t->n.x; i++) {
			float v = (float)vnorm(t->g[0][ind(i, j)].v);
			write_float(fp, v);
		}
	}
	fclose(fp);
	return 0;
}

void init_ball(task_t *t, const double r1, const double r2, const vec2 o, const double mag)
{
	int i, j;
	for (j = 0; j < t->n.y; j++) {
		for (i = 0; i < t->n.x; i++) {
			vec2 p = { t->h.x * i + t->origin.x, t->h.y * j + t->origin.y };
			vec2 d = vsub(p, o);
			double dist = vnorm(d);
			if (dist >= r1 && dist <= r2) {
				double val = mag * impulse((dist - r1) / (r2 - r1));
				t->g[0][ind(i, j)].v = vscal(d, val / dist);
				t->g[0][ind(i, j)].p = val * t->mat.c * t->mat.rho;
			} else {
				t->g[0][ind(i, j)].p = 0.0;
				t->g[0][ind(i, j)].v.x = 0.0;
				t->g[0][ind(i, j)].v.y = 0.0;
			}
		}
	}
}

int save_seismo(const char *file, task_t *t, int first)
{
	int i;
	FILE *fp;
	if (first) {
		fp = fopen(file, "w");
	} else {
		fp = fopen(file, "a");
	}
	if (fp == NULL) {
		perror("Failed to open file");
		return 1;
	}
	for (i = 0; i < t->n.x; i++) {
		fprintf(fp, "%f ", t->g[0][ind(i, t->n.y - 1)].v.y);
	}
	fprintf(fp, "\n");
	fclose(fp);
}


void init_wave(task_t *t, const double w, const vec2 o, const vec2 dir, const double mag)
{
	int i, j;
	for (j = 0; j < t->n.y; j++) {
		for (i = 0; i < t->n.x; i++) {
			vec2 p = { t->h.x * i + t->origin.x, t->h.y * j + t->origin.y };
			double d = vdot(p, dir) - vdot(o, dir);
			if (abs(d) <= w / 2.0) {
				double val = mag * impulse((d + w / 2.0) / w);
				t->g[0][ind(i, j)].v = vscal(dir, val / vnorm(dir));
				t->g[0][ind(i, j)].p = val * t->mat.c * t->mat.rho;
			} else {
				t->g[0][ind(i, j)].p = 0.0;
				t->g[0][ind(i, j)].v.x = 0.0;
				t->g[0][ind(i, j)].v.y = 0.0;
			}
		}
	}
}


